#!/usr/bin/env python
#from __future__ import print_function as fprint
import sys
import os
import hmmer3  #minor modified version of hmmer3
#import MySQLdb
import sqlite3
import numpy as np
import csv
import urllib
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from optparse import OptionParser
from bio.parsers import fasta as fasta_parser
from bio import blosum
#import auto_pdz
#import croc

#12/11/2012
#Will score all the pdz domains from query sequence (if there is more than one)
#change from mysql to sqlite

#11/10 change vpeptide expression to work with situation that
#      peptide len is 5

#known issue:
#query peptide contain 'X' or not AA
#weightM[i][p] will return error

AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
DBUSER = 'root'
DBPASSWD = '19854mdd'
"""
Version: 0.9 1003
fix SQL query problem
!='.' -> not in '.-'

Version: 0.8 0918
Fix hmmsequence, targetsequence
performance increase
add roc

Version: 0.7 0818
Add score normalization into output 

#1. call hmmscan (try to find it from J program) os.system
#2. parse hmmscan file to locate SDRs
#3. for each SDR res, check blosum62 for positive score
#4. (now having the residues from step 3)
#   select () from SDR where blah blah
#    select residues from peptide that bind to the same residues from step 3
#5. build matrix 
#os.system()
"""

def hmmscan_call(hmmdb, fa, eval):
    cmd = 'hmmsearch -E '+str(eval)+' '+hmmdb+' '+fa+'> tmp.hmmscan'
    os.system(cmd)
    hmmscan_out = open('tmp.hmmscan', 'r')
    return hmmscan_out
    
'''
def comb(items, _rs = []):
    #take a list of lists as input, return all possible [combination?]
    #print _tt
    _ts=[]
    _pop = items.pop()
    if type(_pop) is list:
        for i in _pop:
            _ts = [i]+items           #change to also work for str?
            if type(_ts[-1]) is list:
                comb(_ts, _rs)
            else:
                _tt.append(_rs)
    return _rs
'''

def isSDRhit(q_sdrs, d_sdrs, submatrix="blosum62",threhold=0):
    #two dictionaries of sdrs
    #return True is they are similar
    _submatrix = blosum.Matrix(submatrix)
    _isHit = True
    for key in d_sdrs:
        
        #print q_sdrs[key], d_sdrs[key]
        if q_sdrs[key] not in '.-' and d_sdrs[key] not in '.-':
            if _submatrix(q_sdrs[key], d_sdrs[key]) <= threhold:
                _isHit = False
        else:
            _isHit = False
            #print _submatrix(q_sdrs[key], d_sdrs[key])
    #print _isHit
    #raw_input()
            
    return _isHit

#Codes copy from Jonathan's version
class PDZError(Exception):
    pass

class MotifError(Exception):
    pass

def similar_sdrs(c_sdrs, dbname, blosum_matrix):
    """Return list of domain_ids with similar SDRs."""
    M = blosum.Matrix(blosum_matrix)
    con = sqlite3.connect(dbname)
    cur = con.cursor()
    domain_ids = None

    for key in c_sdrs.iterkeys():
        if key.startswith('GLGF'):
            motif = 'GLGF'
            offset = int(key[4:])
        elif key.startswith('GD'):
            motif = 'GD'
            offset = int(key[2:])
        else:
            raise MotifError("Motif needs to be either GLGF or GD.")
        
        if c_sdrs[key] in AminoAcids:
            residues = M.substitutions(c_sdrs[key], 1)
        
            residues = ','.join(['"{0}"'.format(r) for r in residues])
            subquery = """\
            SELECT DISTINCT domainID
            FROM sdr_pdz
            WHERE motif = '{0}'
            AND offset = {1}
            AND residue IN ({2})""".format(motif, offset, residues)
            cur.execute(subquery)
            if domain_ids is None:
                domain_ids = set([row[0] for row in cur])
            else:
                domain_ids = domain_ids.intersection(set([row[0] for row in cur]))

    cur.close()
    con.close()
    return domain_ids

def build_query(peppos, ids):
    """
    MySQL is hopeless with subqueries (they run very slowly); therefore,
    seperate queries are made and the results combined in Python.  Tests show
    this to be much quicker.
    """
    #ids = similar_sdrs(c_sdrs, dbname, blosum_matrix)
    # TO-DO: is peppos counting from C-terminal?  i.e., is peppos = 0
    # the very end?
    sql = """\
    SELECT substr(aa.seq, {0}, 1) AS ss, count(*) AS cc
    FROM (
          SELECT seq
          FROM peptide_pdz
          WHERE domainID IN ({1})
         ) AS aa
    GROUP BY ss"""
    position = 5 - peppos #Please confirm this is right?
    residues = ','.join(["'{0}'".format(x) for x in ids])
    if len(ids) == 0:
        raise PDZError('intersection is empty')
    sql = sql.format(position, residues)
    #print ('SQL', sql)
    return sql
#-------------------------------------------------------------------------------

    

def sql_q(glGf_off, GD_off, i_sequence, peppos, blosumM, j_sequence, dbname, masking = []):

    #return FreqMatrix of ONE position of the peptide
    #i_sequence is the target sequence, j_sequence is the hmm sequence
    
    db = sqlite3.connect(dbname)
    curs = db.cursor()
    curs2 = db.cursor()
    curs.execute("""SELECT distinct s.motif, s.offset
                    FROM sdr as s
                    WHERE s.pep_pos = (?)""", (peppos,))
    motifs_set=curs.fetchall()
    #look up SDR positions from database
    #print motifs_set
    #M = blosum.Matrix(blosumM); #should accept user input for the var
    
    cutoff = 0
    c_sdrs = dict() #compare sdrs | identify SDRs from query sequence
    
    for motif in motifs_set:
        #print motif
        #Create a list of SDR from query sequence
        if motif[0] == "GD":
            try:
                count = abs(motif[1])
                pointer = GD_off
                
                while count != 0:
                    #if there are gaps in between do not count them
                    
                    if motif[1] >=0:
                        pointer += 1
                    else:
                        pointer -= 1
                        
                    if j_sequence[pointer] not in '.-':
                        count -= 1
                        
                sdr = i_sequence[pointer]
                #print sdr
            except (IndexError):
                #print "ERR: SDR not in hmmscan alignment", motif[0], motif[1]
                #if the SDR is not in the alignment
                #not possible to find matches, return ALL 0
                freq = [0] * 20
                return freq
                
            c_sdrs[motif[0]+str(motif[1])] = sdr
            #print c_sdrs
        if motif[0] == "GLGF":
            try:
                count = abs(motif[1])
                pointer = glGf_off
                while count != 0:
                    #if there are gaps in between do not count them
                    
                    if motif[1] >= 0:
                        pointer += 1
                    else:
                        pointer -=1
                    
                    if j_sequence[pointer] not in '-.':
                        count -= 1    
                    
                sdr = i_sequence[pointer]
                #print "sdr:", sdr
            except (IndexError):
                #print "ERR: SDR not in hmmscan alignment", motif[0], motif[1]
                #if the SDR is not in the alignment
                #not possible to find matches, return ALL 0
                freq = [0] * 20
                return freq
            
            c_sdrs[motif[0]+str(motif[1])] = sdr
    #print motif, c_sdrs
    #print "Identification of SDRs from query sequence Done."
    
    results = []

    #---------CODE FROM JJE---------#
    #Looking for similar residues from database peptides
    freq = [0] * 20
    #ids = domain IDs, which have similar SDRs to the query domain.
    ids_2 = similar_sdrs(c_sdrs, dbname, blosumM)
    ids = []
    #MASKING AT HERE, ids <-- check!
    for i in ids_2:
        if i not in masking:
           ids.append(i) 
    
    try:
        sql = build_query(peppos, ids)
    except PDZError:
        # There were no domains in the intersection. (so the IN clause is empty
        # and therefore an SQL error).
        return freq

    con = sqlite3.connect(dbname)
    cur = con.cursor()
    try:
        cur.execute(sql)
        for row in cur:
            try:
                idx = AminoAcids.index(row[0])
            except ValueError:
                pass
            else:
                freq[idx] = int(row[1])
    finally:
        cur.close()
        con.close()
    #print 'Freq:', freq
    return freq
    
#COPY and EDIT from Jonathan's Predikin3, kinase.py
#TODO: Create backgroud freq from database?
def create_weights(frequencies):
    #frequencies = self.frequencies
    #priors = self._get_priors()
    priors = {'PDZ':   {'A': 0.069, 'C': 0.017, 'D': 0.054, 'E': 0.069,
                       'F': 0.034, 'G': 0.071, 'H': 0.025, 'I': 0.048,
                       'K': 0.051, 'L': 0.092, 'M': 0.021, 'N': 0.038,
                       'P': 0.066, 'Q': 0.045, 'R': 0.060, 'S': 0.086,
                       'T': 0.054, 'V': 0.064, 'W': 0.010, 'Y': 0.027},
              }
    prior = priors['PDZ'] #using 'PDZ' at the moment
    
    #weights = np.zeros([5, 20])
    weights = {}
    sqrt = np.sqrt
    log2 = np.log2

    for i, row in enumerate(frequencies):
        weights[i] = {}
        if i == 1:
            for j,freq in enumerate(row):
                weights[i][AminoAcids[j]] = 0
        else:
            N = sum(row)
            if N > 0:
                for j,freq in enumerate(row):
                    #print j, freq
                    p = prior[AminoAcids[j]]
                    weight = log2(((freq + (sqrt(N) / 20.0)) / (N + sqrt(N))) / p)
                    weights[i][AminoAcids[j]] = weight
            else:
                return False
    return weights

def is_valid(M):
    # 0=invalid, 1=suspicious, 2=valid
    r = 'VALID'
    for row in M:
        if sum(row) == 0:
            r = 'INVALID'
            break
        elif sum(row) < 5:
            r = 'SUSPICIOUS'
            break
        
    if r == 'SUSPICIOUS' and sum(M[0]) > 5 and sum(M[2]) > 5:
        r = r + ' *'
    return r

def score_peptide(peptide, weightM, roc_scores = [], isbind = -1):
    score = 0
    minsum = 0
    maxsum = 0
    
    for i, p in enumerate(peptide[::-1][:5]):
        skip = False
        #print peptide
        #print peptide[::-1][:5]
        #print i, p
        #if len(peptide) != 5:
        #    raw_input()
        try:
            raw_score = weightM[i][p]
        except:
            skip = True
            #raw_score = float(sum(weightM[i].itervalues())/ len(weightM[i]))
            #print "Len:", len(weightM[i]), raw_score
        #normalize the score
        if not skip:
            minsum += min(weightM[i].itervalues())
            
            maxsum += max(weightM[i].itervalues())
            score += raw_score
        
    
    nscore = (score-minsum)/(maxsum-minsum)
    if nscore < 0:
        print "minsum:", minsum
        print "maxsum:", maxsum
        print "score:", score
        print peptide[::-1][:5]
    
    if isbind == -1:
        return nscore
    else:
        print nscore, isbind
        roc_scores.add(nscore, isbind)
    
        return roc_scores

def data2file(scores, dbname):
    i = dbname[-1]
    outfile = 'datafile{0}.dat'.format(i)
    data = scores.score_labels
    with open (outfile, 'w') as fout:
        for k in data:
            values = data[k]
            for v in values:
                writer = csv.writer(fout, delimiter = ' ', lineterminator = '\n')
                writer.writerow([k, v])




def getPWM(query_seq, dbname, masking = []):
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    tmpPath = "scanHmmer.tmp"
    tmpFile = open(tmpPath, "w" )
        #####################################################################
    #   variables
    ##############
    #import csv table for validation

    hmmscan_thres = 0.01
    hmmdb = 'strucaln.hmm'
    options_blosum = "blosum62"
    options_length = 5
    #####################################################################
    previous_protein = ''
    previous_pdznum = ''
    
    
    
    #pdz_record = SeqIO.parse(infile, "fasta").next()
    
    SeqIO.write(query_seq, tmpFile, "fasta")
    tmpFile.close()
    read_data = hmmscan_call(hmmdb, tmpPath, hmmscan_thres)
    #raise
    #cnt_record = 0
    #cnt_hit = 0
    #cnt_domain = 0
    pwms = dict()
    for record in hmmer3.parse(read_data):
        #cnt_record += 1
        for hit in record:
            #cnt_hit += 1
            for domain in hit:
                #cnt_domain += 1
                
                glGf_off = domain.hmm_sequence.find('glGf') + 3
                #index of glGf motif
                GD_off = domain.hmm_sequence.find('GD') + 1
                #index of GD motif
                #load the fasta file into fasta parser
                freqM = []   #output frequency matrix
                weightM = [] #output PWM matrix
                for p_position in range(options_length):
                    #careful with hmm_sequence and target_sequence
                    #print hit.name, p_position
                    
                    freqM.append(sql_q(glGf_off, GD_off ,
                                       domain.target_sequence,
                                       p_position,
                                       options_blosum,
                                       domain.hmm_sequence, dbname))
                    #print freqM
                    weightM = create_weights(freqM)
                '''
                #=====OUTPUT=====
                #print weightM
                #print "Weight:", weightM
                if weightM:
                    #print "weight matrix ok!"
                    #x = is_valid(freqM)
                    score = 0
                    #print "len pep:", len(vpeptide)
                    if len(vpeptide) >= 4:
                        roc_scores = score_peptide(vpeptide, weightM, roc_scores, isbind)
                '''
                pwms[(domain.ali_from, domain.ali_to)] = weightM
            try:
                os.unlink( tmpPath )
            except OSError:
                pass
            
            #print roc_scores
            #print '\t'.join([record.query, hit.name, domain.id, x])
            
    #data2file(roc_scores, dbname)
            #return weightM
    #print cnt_record, cnt_hit, cnt_domain
    if len(pwms) > 0:
        return pwms
    else:
        return False


    
    
def main(validtable, dbname, masking = []):
    #masking is a list of Uniprot ID, which will not be used
    #in the training data
    AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')
    #auto_pdz.main()
    '''
    usage = "usage: %prog [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-i", "--input",
                      dest="inputfasta", help="input fasta sequence file")
    parser.add_option("-d", "--hmmdb", dest="hmmdb", help="hmm database")
    parser.add_option("-b", "--blosum", dest="blosum", default="blosum62",
                      help="choose blosum matrix [default: %default]")
    parser.add_option("-E", dest="hmmthreshold", default =0.01,
                      help="Threshold of E-value of hmmscan [default: %default]")
    parser.add_option("-l", "--length", dest="length", default = 5,
                      help="length of predicting peptide (max=5 atm) [default: %default]")
        
    (options, args) = parser.parse_args()

    hmmscan_thres = options.hmmthreshold
    # threshold set to 0 to ignore the check for testing
    '''
    #####################################################################
    #   variables
    ##############
    #import csv table for validation
    raw_validation_table = csv.reader(open('database/'+validtable))
    validation_table = sorted(raw_validation_table)
    hmmscan_thres = 0.01
    hmmdb = 'strucaln.hmm'
    options_blosum = "blosum62"
    options_length = 5
    #####################################################################
    previous_protein = ''
    previous_pdznum = ''
    roc_scores = croc.ScoredData()
    
    for row in validation_table:
        #uniprot_acc, start, end, peptide, binding [0 nonbinder]
        #pdz_num = row[3]
        pdz_loc = (int(row[1]), int(row[2]))
        inputfasta = row[0]+".fasta"
        
        try:
            seq_record = SeqIO.parse("validate_data/"+inputfasta, "fasta").next()
        except(IOError):
            print inputfasta, "not found. Retrieving Uniprot file online ... "
            urllib.urlretrieve("http://www.uniprot.org/uniprot/" + inputfasta,
                               "validate_data/" + inputfasta)
            seq_record = SeqIO.parse("validate_data/"+inputfasta, "fasta").next()
        
        if pdz_loc != (0, 0):
            pdz_record = seq_record[pdz_loc[0]-1:pdz_loc[1]-1]
        tmpPath = "scanHmmer.tmp"
        tmpFile = open(tmpPath, "w" )
        #print "seq:", seq
        #ss = Seq(seq)
        #record = SeqRecord(seq, acc, "", "")
        #if pdz_record.id == 'sp|Q9HAP6|LIN7B_HUMAN':
        #    print pdz_record
        #    raw_input('wait')
        SeqIO.write(pdz_record, tmpFile, "fasta")
        
        tmpFile.close()
        
        
        vpeptide = row[3]
        #print "row4:", row[4]
        #print row
        isbind = int(row[4])

            
        #acclerrate only when the input file are sorted by the domain ID
        #by not calculating the freqM if the same PDZ domain is involved
        #in the prediction.
        if previous_protein == inputfasta and previous_pdznum == pdz_loc:
            if weightM:
                #x = is_valid(freqM)
                score = 0
                #print vpeptide
                if len(vpeptide) >= 4:
                    roc_scores = score_peptide(vpeptide, weightM, roc_scores, isbind)

                #else:
                #    print 'not enough data'
            previous_pdznum = pdz_loc
            previous_protein = inputfasta
        else:
            #if previous_protein != inputfasta:
            #print 'Call hmmsearch'
            read_data = hmmscan_call(hmmdb, tmpPath, hmmscan_thres)
            #print 'call done'
            #input_fasta = fasta_parser.FastaFile(options.inputfasta)
            for record in hmmer3.parse(read_data):
                for hit in record:
                    for domain in hit:
                        #if domain.id == pdz_num:
                            
                            
                            #print domain.id
                            
                            glGf_off = domain.hmm_sequence.find('glGf') + 3
                            #index of glGf motif
                            GD_off = domain.hmm_sequence.find('GD') + 1
                            #index of GD motif
                            #find SDRx
                            #load the fasta file into fasta parser
                            freqM = []   #output frequency matrix
                            #print freqM
                            weightM = [] #output PWM matrix
                            for p_position in range(options_length):
                                #careful with hmm_sequence and target_sequence
                                #print hit.name, p_position
                                
                                freqM.append(sql_q(glGf_off, GD_off ,
                                                   domain.target_sequence,
                                                   p_position,
                                                   options_blosum,
                                                   domain.hmm_sequence, dbname))
                                #print freqM
                                weightM = create_weights(freqM)
            
                            #=====OUTPUT=====
                            #print weightM
                            #print "Weight:", weightM
                            if weightM:
                                #print "weight matrix ok!"
                                #x = is_valid(freqM)
                                score = 0
                                #print "len pep:", len(vpeptide)
                                if len(vpeptide) >= 4:
                                    
                                    roc_scores = score_peptide(vpeptide, weightM, roc_scores, isbind)

                            #else:
                            #    print 'not enough data'
            previous_pdznum = pdz_loc
            previous_protein = inputfasta
            try:
                os.unlink( tmpPath )
            except OSError:
                pass
            #print roc_scores
            #print '\t'.join([record.query, hit.name, domain.id, x])
            
    #data2file(roc_scores, dbname)
    try:
        roc = croc.ROC(roc_scores.sweep_threshold())
        return roc, roc.area()
    except Exception, e:
        print e
        return False, False

def seq2seqrecord(seq):
    #convert query sequence to SeqRecord format
    record = SeqRecord(Seq(seq,IUPAC.protein), id = "Query Sequence", name="Query",
                       description = "Query Sequence")
    
    return record

def fasta2seqrecord(infile):
    #convert fasta file to SeqRecord format
    
    record = SeqIO.parse(infile, "fasta").next()
    
    return record

def scoring(query_sequence, peptides, dbname):
    pwms = getPWM(query_sequence, dbname, [])

    results = dict()
    if pwms is not False:
        for k in pwms:
            
            scores = []
            if pwms[k] is not False:
                for peptide in peptides:
                    if len(peptide) == 5:
                        scores.append(score_peptide(peptide, pwms[k]))
                    else:
                        print "Peptide length is not 5"
                        scores.append(0)
            results[k] = scores
        return results, len(pwms)
    else:
        return False

if __name__ == '__main__':
    #AminoAcids = list('ACDEFGHIKLMNPQRSTVWY')

    q_seq = "TVEIIRELSDALGISIAGGKGSPLGDIPIFIAMIQANGVAARTQKLKVGDRIVSINGQPLDGLSHTDAVNLLKNAFGRIILQVVADT"
    #seq = fasta2seqrecord("test.fasta")
    seq = seq2seqrecord(q_seq)
    peptides = ["TKVLV","KKTTV"]
    scores, pdz_count = scoring (seq, peptides, "pdzvalid.db")
    print pdz_count, scores
    raise
    '''
    validtable = 'dummy_test.dat'
                        #a, b = main(validtable, dbname);
    for i in xrange(10):
        #testfile = 
        dbname = 'pdzvalid_{0}'.format(i + 1)
        roc, auc = main(validtable, dbname)
    '''
    #test(','/database/pdzvalid', [])
    '''
    for dbnum in range(0, 10):
        dbname = 'pdzvalid_' + str(dbnum + 1)
        validtable = 'test_group'+str(dbnum)+'.dat'
        a, b = main(validtable, dbname);
        print a 
        print b
    '''

